package cn.imatu.framework.log.utils;

import brave.Tracer;
import brave.Tracing;
import brave.propagation.TraceContext;
import cn.imatu.framework.core.utils.spring.SpringUtils;
import cn.imatu.framework.log.common.LogCons;
import org.apache.commons.lang3.StringUtils;

import java.util.Objects;
import java.util.function.BiConsumer;

/**
 * @author shenguangyang
 */
public class BraveTracerUtils {

    public static void addTraceId(BiConsumer<String, String> headers) {
        Tracer tracer = Tracing.currentTracer();
        if (Objects.nonNull(tracer)) {
            TraceContext context = tracer.currentSpan().context();
            headers.accept(LogCons.TRACE_ID_HEADER, context.traceIdString());
            headers.accept(LogCons.SPAN_ID_HEADER, context.spanIdString());
            if (StringUtils.isBlank(context.parentIdString())) {
                headers.accept(LogCons.PARENT_SPAN_ID_HEADER, context.spanIdString());
            } else {
                headers.accept(LogCons.PARENT_SPAN_ID_HEADER, context.parentIdString());
            }
        }
    }

    public static String getTraceString() {
        Tracer tracer = SpringUtils.getBean(Tracer.class);
        if (Objects.isNull(tracer) || Objects.isNull(tracer.currentSpan()) || Objects.isNull(tracer.currentSpan().context())) {
            return "";
        }
        return tracer.currentSpan().context().traceIdString();
    }


}
